# -*- coding: utf-8 -*-
"""
翻转
"""

import os
import sys
import argparse
import random
import voc_xml
import cv2
from voc_xml import CreateXML


def flip_img(src, flip_type):
    """
    翻转图像
    Args:
        src:输入图像
        flip_type:翻转类型，1水平翻转，0垂直翻转，-1水平垂直翻转
    return:
        fliped_img:翻转后的图像

    """
    fliped_img = cv2.flip(src, flip_type)
    return fliped_img


def flip_xy(x, y, img_w, img_h, flip_type):
    """
    翻转坐标点
    Args:
        x,y：坐标点
        imgw,imgh:翻转图像宽高
        flip_type:翻转类型，1水平翻转，0垂直翻转，-1水平垂直翻转
    return:
        fliped_x,fliped_y:翻转后坐标 
    """
    if 1 == flip_type:
        fliped_x = img_w - x
        fliped_y = y
    elif 0 == flip_type:
        fliped_x = x
        fliped_y = img_h - y
    elif -1 == flip_type:
        fliped_x = img_w - x
        fliped_y = img_h - y
    else:
        print('flip type err')
        return
    return fliped_x, fliped_y


def flip_box(box, img_w, img_h, flip_type):
    """
    翻转目标框
    Args:
        box:目标框坐标[xmin,ymin,xmax,ymax]
        imgw,imgh:图像宽高
        flip_type:翻转类型,1水平翻转，0垂直翻转，-1水平垂直翻转
    return:
        fliped_box:翻转后的目标框
    """
    x1, y1 = flip_xy(box[0], box[1], img_w, img_h, flip_type)
    x2, y2 = flip_xy(box[2], box[3], img_w, img_h, flip_type)
    xmin, xmax = min(x1, x2), max(x1, x2)
    ymin, ymax = min(y1, y2), max(y1, y2)
    fliped_box = [xmin, ymin, xmax, ymax]
    return fliped_box


def flip_xml(flip_img_name, xml_tree, flip_type):
    """
    翻转xml
    Args:
        flip_img_name:翻转后图片保存名
        xml_tree:待翻转的xml ET.parse()
        flip_type:翻转类型,1水平翻转，0垂直翻转，-1水平垂直翻转
    return:
        createdxml : 创建的xml CreateXML对象   
    """
    root = xml_tree.getroot()
    size = root.find('size')
    img_w, img_h, depth = int(size.find('width').text),\
                          int(size.find('height').text),\
                          int(size.find('depth').text)
    createdxml = CreateXML(flip_img_name, int(img_w), int(img_h), depth)
    for obj in root.iter('object'):
        obj_name = obj.find('name').text
        xml_box = obj.find('bndbox')
        xmin = int(xml_box.find('xmin').text)
        ymin = int(xml_box.find('ymin').text)
        xmax = int(xml_box.find('xmax').text)
        ymax = int(xml_box.find('ymax').text)
        box = flip_box([xmin, ymin, xmax, ymax], img_w, img_h, flip_type)
        if (box[0] >= box[2]) or (box[1] >= box[3]):
            continue
        createdxml.add_object_node(obj_name, box[0], box[1], box[2], box[3])
    return createdxml
    

def flip_img_xml(img, xml_tree, flip_img_name, flip_type):

    """
    翻转图像和xml目标框
    Args:
        img：源图像
        xml_tree：待crop的xml ET.parse()
        crop_img_name:翻转图片命名
        flip_type:翻转类型
    return:
        fliped_img,fliped_xml : 裁剪完成的图像和xml文件        
    """
    fliped_img = flip_img(img, flip_type)
    fliped_xml = flip_xml(flip_img_name, xml_tree, flip_type)
    return fliped_img, fliped_xml


def batch_flip_dir(imgs_dir, xmls_dir, imgs_save_dir, xmls_save_dir, img_suffix, name_suffix,
                   flip_types, random_flip=False):
    """
    翻转指定路径下所有图片和xml
    Args:
        imgs_dir,xmls_dir:待翻转图片和xml路径
        imgs_save_dir,xmls_save_dir：图片和xml保存路径
        img_suffix:图片可能的后缀名['.jpg','.png','.bmp',..]
        name_suffix: 处理完成的图片、xml的命名标识
        flip_types: 每张图执行的翻转类型[type1,type2,...],翻转类型共三种,1水平翻转，0垂直翻转，-1水平垂直翻转
        random_flip:是否随机选择翻转类型,与flip_type互斥     
    """
    for root, dirs, files in os.walk(xmls_dir):
        for xml_name in files:
            xml_file = os.path.join(xmls_dir, xml_name)
            img_file = None
            for suffix in img_suffix:
                if os.path.exists(os.path.join(imgs_dir, xml_name.split('.')[0]+suffix)):
                    img_file = os.path.join(imgs_dir, xml_name.split('.')[0]+suffix)
                    break
            if img_file is None:
                print("image for \'{}\' does not exists", xml_name)
                continue
            img = cv2.imread(img_file)
            types = flip_types
            if random_flip:
                types = [random.randint(-1, 1)]
            for flip_type in types:
                flip_img_name = "{}_{}_type-{}_.{}".format(xml_name.split('.')[0],
                                                           name_suffix, flip_type,
                                                           img_file.split('.')[-1])
                img_flip, xml_flip = \
                    flip_img_xml(img, voc_xml.get_xml_tree(xml_file), flip_img_name, flip_type)
                cv2.imwrite(os.path.join(imgs_save_dir, flip_img_name), img_flip)
                xml_flip.save_xml(xmls_save_dir, flip_img_name.split('.')[0]+'.xml')


def start_flip(root_dir, save_dir):
    img_dir = os.path.join(root_dir, "JPEGImages")
    xml_dir = os.path.join(root_dir, "Annotations")
    flip_img_save_dir = os.path.join(save_dir, "JPEGImages")
    flip_xml_save_dir = os.path.join(save_dir, "Annotations")
    if not os.path.exists(flip_img_save_dir):
        os.makedirs(flip_img_save_dir)
    if not os.path.exists(flip_xml_save_dir):
        os.makedirs(flip_xml_save_dir)

    img_suffix = ['.jpg', '.png', '.bmp']

    # add to file name to identify flip methods:
    # mirror, vertical or mirror_vertical
    name_suffix = 'flip'

    # 1 - mirror; 0 - vertical; -1 - mirror_vertical
    flip_types = [1, 0, -1]

    # random flip
    # mutual exclusion with flip_types
    random_flip = True

    batch_flip_dir(img_dir, xml_dir, flip_img_save_dir,
                   flip_xml_save_dir, img_suffix, name_suffix,
                   flip_types, random_flip)


if __name__ == '__main__':
    try:
        ROOT_DIR = sys.argv[1]
        SAVE_DIR = sys.argv[2]
    except:
        PARSER = argparse.ArgumentParser(description="arguments")
        PARSER.add_argument("--root_dir", type=str,
                            default="/home/chujie/PycharmProjects/data_tool/data/voc2012",
                            help="VOC dataset root directory")
        PARSER.add_argument("--save_dir", type=str,
                            default=os.path.join(os.getenv("HOME"),
                                                 "PycharmProjects/data_tool/data/flip_voc"),
                            help="result save dir")
        ARGS = PARSER.parse_args()
        ROOT_DIR = ARGS.root_dir
        SAVE_DIR = ARGS.save_dir
    start_flip(ROOT_DIR, SAVE_DIR)
